#include "tppmc/TF1J.h"

#include <TRandom.h>
#include <TRandom3.h>
#include <TF1.h>
#include <TMath.h>

#include <unistd.h>
#include <iostream>
#include <fstream>
#include <stdlib.h>
#include <time.h>
#include <math.h>
#include <time.h>

using namespace std;
#define QUARK_TYPES 9  
#define N_PART_TYPES 9*QUARK_TYPES //pion, kaon, proton with +/-/0 charges of each
#define QS_NBINS 230

enum PartonType{GLUON,UQUARK,UBAR,DQUARK,DBAR,SQUARK,SBAR,CQUARK,CBAR,BQUARK,BBAR};

TF1J::TF1J()
: TF1()
, fAlpha(NULL)
, fBeta(NULL)
, fGamma(NULL)
{
    // allocate memory for our cache
    fIntegral_cache = new double**[N_PART_TYPES];
    r1_cache = new double**[N_PART_TYPES];
    total_cache = new double*[N_PART_TYPES];
    
    for (int i = 0; i < N_PART_TYPES; ++i) {
        fIntegral_cache[i] = NULL;
        fIntegral_cache[i] = new double*[QS_NBINS];
        
        r1_cache[i] = NULL;
        r1_cache[i] = new double*[QS_NBINS];
        
        total_cache[i] = NULL;
        total_cache[i] = new double[QS_NBINS];
    }
    
    for(int i=0; i<N_PART_TYPES; i++){
        for(int j=0; j<QS_NBINS; j++){
            fIntegral_cache[i][j] = NULL;
            r1_cache[i][j] = NULL;
        }
    }
    //fIntegral = NULL; 
}

TF1J::TF1J(const char *name,Double_t (*fcn)(Double_t *, Double_t *), Double_t xmin, Double_t xmax, Int_t npar)
: TF1(name, fcn, xmin, xmax, npar)
, fAlpha(NULL)
, fBeta(NULL)
, fGamma(NULL)
{
   // allocate memory for our cache
   fIntegral_cache = new double**[N_PART_TYPES];
   r1_cache = new double**[N_PART_TYPES];
   total_cache = new double*[N_PART_TYPES];
   
   for (int i = 0; i < N_PART_TYPES; ++i) {
      fIntegral_cache[i] = NULL;
      fIntegral_cache[i] = new double*[QS_NBINS];
      
      r1_cache[i] = NULL;
      r1_cache[i] = new double*[QS_NBINS];
      
      total_cache[i] = NULL;
      total_cache[i] = new double[QS_NBINS];
   }
   
   for(int i=0; i<N_PART_TYPES; i++){
      for(int j=0; j<QS_NBINS; j++){
         fIntegral_cache[i][j] = NULL;
         r1_cache[i][j] = NULL;
      }
   }
   //fIntegral = NULL; 
}

TF1J::~TF1J()
{
   for (int i = 0; i < N_PART_TYPES; ++i) {
      for (int j = 0; j < QS_NBINS; ++j){
         if(fIntegral_cache[i][j]){ delete [] fIntegral_cache[i][j]; }
         if(r1_cache[i][j]){ delete [] r1_cache[i][j]; }
      }
      if(fIntegral_cache[i]){ delete [] fIntegral_cache[i]; }
      if(r1_cache[i]){ delete [] r1_cache[i]; }
      if(total_cache[i]){ delete [] total_cache[i]; }
   }
   if(fIntegral_cache){ delete [] fIntegral_cache; }
   if(r1_cache){ delete [] r1_cache; }
   if(total_cache){ delete [] total_cache; }
   
   if(fAlpha){ delete [] fAlpha; }
   if(fBeta){ delete [] fBeta; }
   if(fGamma){ delete [] fGamma; }
}


Double_t TF1J::GetRandom(Double_t xmin, Double_t xmax)
{
//   std::cout << "TF1J::GetRandom("<<xmin<<", "<<xmax<<")"<<std::endl;
   qs = pow(GetParameter(1),2);
   if(qs>100000){
      cout << "Q^2 > 100k, setting Q^2 = 100k (TF1J)" << endl;
      qs=100000;
   }
   type_j = static_cast<int>(GetParameter(4))-1; //PARTICLE SPECIES 0=pion, 1=kaon, 2=proton
   if(type_j!=0 && type_j!=1 && type_j!=2){ 
      cout << "particle_id != 0 [pion], 1 [kaon], or 2 [proton], TF1J wont segfault, so I'm exiting the program" << endl;
      exit(-1);
   }
   
   int chg_j = static_cast<int>(GetParameter(5))+1;
   if(chg_j!=0 && chg_j!=1 && chg_j!=2){ 
      cout << "invalid chg_j, TF1J wont segfault, so I'm exiting the program" << endl;
      exit(-1);
   }
   
   int qtyp = static_cast<int>(GetParameter(2));
   // c=cbar and b=bbar in DSS, no need to store duplicates, so compress the index
   // (takes us down from 11 to 9)
   if(qtyp==CBAR)
      qtyp = CQUARK; 
   else if( (qtyp==BQUARK) || (qtyp==BBAR) )
      qtyp = CBAR;
   
//   std::cout << "Q2 = " << qs << " hadron = " << type_j << " charge = " <<
//   chg_j << " parton = " << qtyp << std::endl;
   
   // set the index based on hadron type, charge and quark
   
   type_j = type_j*(3*QUARK_TYPES)+(chg_j*QUARK_TYPES)+qtyp; 
   
   if(type_j>(N_PART_TYPES-1)){
      cout << "woah woah woah! type_j too large in TF1J.cc, this is going to get messy..." << endl;
      cout << static_cast<int>(GetParameter(4))-1 << " " << chg_j << " " << qtyp << " " << type_j << endl;
   }
   
   
   ////////////////////////// WHICH BIN DOES THIS Q^2 VALUE BELONG TO //////////////////////////
   double macro_bin_edges[] = {1, 1.25, 1.5, 2.5, 4, 6.4, 10, 15, 25, 40, 64, 100, 180, 580, 1000, 1800, 3200, 5800, 10000, 18000, 32000, 58000, 100000};
   for(int i4=0; i4<QS_NBINS; i4++){
      mod_fact = i4%10;
      lower = floor(static_cast<double>(i4)/10.0);
      micro_bin_edge = macro_bin_edges[lower+1]-macro_bin_edges[lower];
      micro_bin_edge /= 10.0;
      micro_bin_edge *= mod_fact;
      micro_bin_edge += macro_bin_edges[lower];
      
      //cout << micro_bin_edge << endl;
      if(qs < micro_bin_edge){ 
         qs_j = i4-1;
         break;
      }
      if(qs == micro_bin_edge){
         qs_j = i4;
         break;
      }
   }
//   cout << "qs_j: " << qs_j << " qs: " << qs << endl;
   if(qs_j>QS_NBINS){ cout << "____________________QS_J > QS_NBINS, GOING TO SEGFAULT (TF1J.cc)__________________" << endl; }
   if(qs_j<0){ cout << "_____________________QS_J < 0, GOING TO SEGFAULT (TF1J.cc)________________" << endl; }
   
   
   Double_t dx = (fXmax-fXmin)/fNpx;
   if(!fAlpha){ fAlpha = new Double_t[fNpx]; }
   if(!fBeta){ fBeta = new Double_t[fNpx]; }
   if(!fGamma){ fGamma = new Double_t[fNpx]; }
   
//   cout << "fIntegral_cache  [" << type_j << "][" << qs_j << "]: " << fIntegral_cache[type_j][qs_j] << endl;
   //////////////////////////// Is this type_j and qs_j combination already cached? if not, build it /////////////////////
   if (fIntegral_cache[type_j][qs_j]==NULL) {
      fIntegral_cache[type_j][qs_j] = new Double_t[fNpx+1];
      fIntegral_cache[type_j][qs_j][0] = 0.0;
      
      r1_cache[type_j][qs_j] = new Double_t[fNpx];
//      cout << "alloced properly?[" << type_j << "][" << qs_j << "]: " << fIntegral_cache[type_j][qs_j] << endl;
      
      Double_t integ;
      Int_t intNegative = 0;
      Int_t i;
      for (i=0;i<fNpx;i++) {
         integ = Integral(Double_t(fXmin+i*dx), Double_t(fXmin+i*dx+dx));
         if (integ < 0) {intNegative++; integ = -integ;}
         fIntegral_cache[type_j][qs_j][i+1] = fIntegral_cache[type_j][qs_j][i] + integ;
      }
      if (intNegative > 0) {
         Warning("GetRandom","function:%s has %d negative values: abs assumed",GetName(),intNegative);
      }
      if (fIntegral_cache[type_j][qs_j][fNpx] == 0) {
         Error("GetRandom","Integral of function is zero");
         return 0;
      }
      
      total_cache[type_j][qs_j] = fIntegral_cache[type_j][qs_j][fNpx];
      
      for (i=1;i<=fNpx;i++) {  // normalize integral to 1
         fIntegral_cache[type_j][qs_j][i] /= total_cache[type_j][qs_j];
      }
      
      
      for (i=0;i<fNpx;i++) {
         x0 = fXmin+i*dx;
         r1_cache[type_j][qs_j][i] = Integral(x0,x0+0.5*dx)/total_cache[type_j][qs_j];
      }
      
   }
   //////////////////////////// Set temp array (fIntegral) to be used by BinarySearch //////////////////
   
   fIntegral = fIntegral_cache[type_j][qs_j];
   
   for (i=0;i<fNpx;i++) {
      x0 = fXmin+i*dx;
      r2 = fIntegral[i+1] - fIntegral[i];
      r1 = r1_cache[type_j][qs_j][i];
      r3 = 2*r2 - 4*r1;
      if (TMath::Abs(r3) > 1e-8) fGamma[i] = r3/(dx*dx);
      else           fGamma[i] = 0;
      fBeta[i]  = r2/dx - fGamma[i]*dx;
      fAlpha[i] = x0;
      fGamma[i] *= 2;
   }
   
   
   // return random number
   Int_t nbinmin = static_cast<Int_t>((xmin-fXmin)/dx);
   Int_t nbinmax = static_cast<Int_t>((xmax-fXmin)/dx)+2;
   if(nbinmax>fNpx) nbinmax=fNpx;
   
   Double_t pmin=fIntegral[nbinmin];
   Double_t pmax=fIntegral[nbinmax];
   
   Double_t r,x,xx,rr;
   do {
      r  = gRandom->Uniform(pmin,pmax);
      
      Int_t bin  = TMath::BinarySearch(fNpx,fIntegral,r);
      rr = r - fIntegral[bin];
      
      if(fGamma[bin] != 0)
         xx = (-fBeta[bin] + TMath::Sqrt(fBeta[bin]*fBeta[bin]+2*fGamma[bin]*rr))/fGamma[bin];
      else
         xx = rr/fBeta[bin];
      x = fAlpha[bin] + xx;
   } while(x<xmin || x>xmax);
   
   return x;
}
